module Testnet.TestEnumGenerator
  ( genTestType
  , genAllConstructorsList
  ) where

import           Language.Haskell.TH (Body (NormalB), Con (..), Dec (DataD, ValD),
                   Exp (ConE, ListE), Info (TyConI), Name, Pat (VarP), Q, mkName, nameBase, reify)

-- | Create a datatype with the same constructors as the given type, but with a "Test" prefix and no arguments.
-- For example, if the input type is 'Maybe', the output type will be 'TestMaybe', with constructors 'TestNothing'
-- and 'TestJust', but 'TestJust' will have no arguments.
genTestType :: Name -> Q [Dec]
genTestType typeName = do
  TyConI (DataD _ _ _ _ constructors _) <- reify typeName
  let newConstructors = map (makeSimpleConstructor . addTestPrefix) $ concatMap getConstructorName constructors
  return [DataD [] (addTestPrefix typeName) [] Nothing newConstructors []]
 where
  addTestPrefix :: Name -> Name
  addTestPrefix name = mkName $ "Test" ++ nameBase name

  makeSimpleConstructor :: Name -> Con
  makeSimpleConstructor n = NormalC n []

-- | Generate a declaration with a list of all constructors of a type. For example, if the input type is 'Maybe',
-- the output will be 'allMaybeConstructors = [Nothing, Just]'. Obviously, this will only work if all constructors
-- have types that can be unified, like is the case with nullary constructors like the ones generated by 'genNewType'.
genAllConstructorsList :: Name -> Q [Dec]
genAllConstructorsList typeName = do
  constructorList <- getAllConstructors typeName
  return [ValD (VarP $ mkName $ "all" ++ nameBase typeName ++ "Constructors") (NormalB constructorList) []]
 where
  getAllConstructors :: Name -> Q Exp
  getAllConstructors typeName' = do
    TyConI (DataD _ _ _ _ constructors _) <- reify typeName'
    return $ ListE $ map ConE $ concatMap getConstructorName constructors

-- | Obtain the name or names from a constructor
getConstructorName :: Con -> [Name]
getConstructorName (NormalC na _) = [na]
getConstructorName (RecC na _) = [na]
getConstructorName (InfixC _ na _) = [na]
getConstructorName (ForallC _ _ con) = getConstructorName con
getConstructorName (GadtC nas _ _) = nas
getConstructorName (RecGadtC nas _ _) = nas
